# System Import

# ThirdParty Import
import taichi as Ti
# Local Import
import Math
import Config.GlobalParameters as GP


# General Jacobi Kernel
@Ti.kernel
def Poisson(InputBuffer: Ti.template(), OutputBuffer: Ti.template(), DivBuffer: Ti.template()):
    for i, j in InputBuffer:
        L = Math.Sample(InputBuffer, i - 1, j)
        R = Math.Sample(InputBuffer, i + 1, j)
        B = Math.Sample(InputBuffer, i, j - 1)
        T = Math.Sample(InputBuffer, i, j + 1)
        Div = DivBuffer[i, j]
        OutputBuffer[i, j] = (L + R + B + T - Div * GP.DeltaX2) * 0.25


# Specific Jacobi Kernel
@Ti.kernel
def Diffuse(InputBuffer: Ti.template(), OutputBuffer: Ti.template()):
    for i, j in InputBuffer:
        L = Math.Sample(InputBuffer, i - 1, j)
        R = Math.Sample(InputBuffer, i + 1, j)
        B = Math.Sample(InputBuffer, i, j - 1)
        T = Math.Sample(InputBuffer, i, j + 1)
        W2 = InputBuffer[i, j]
        OutputBuffer[i, j] = (
                                     (L + R + B + T) * GP.DeltaTime * GP.ViscosityFactor + W2 * GP.DeltaX2) \
                             / \
                             (GP.DeltaX2 + 4 * GP.DeltaTime * GP.ViscosityFactor)


# Scalar Diffuse
@Ti.kernel
def ScalarDiffuse(InputBuffer: Ti.template(), OutputBuffer: Ti.template(), DiffuseFactor: float):
    for i, j in InputBuffer:
        L = Math.Sample(InputBuffer, i - 1, j)
        R = Math.Sample(InputBuffer, i + 1, j)
        B = Math.Sample(InputBuffer, i, j - 1)
        T = Math.Sample(InputBuffer, i, j + 1)
        W = InputBuffer[i, j]
        OutputBuffer[i, j] = (L + R + B + T - 4 * W) * GP.DeltaTime * DiffuseFactor * GP.InvDeltaX2 + W


# TempBuffer
TempScalarBuffer1 = Ti.field(float, shape=(GP.FluidsResolution, GP.FluidsResolution))
TempScalarBuffer2 = Ti.field(float, shape=(GP.FluidsResolution, GP.FluidsResolution))
TempVectorBuffer1 = Ti.Vector.field(2, float, shape=(GP.FluidsResolution, GP.FluidsResolution))
TempVectorBuffer2 = Ti.Vector.field(2, float, shape=(GP.FluidsResolution, GP.FluidsResolution))


# Compact Poisson
@Ti.func
def Reflection(Index):
    IndexTemp = abs(Index)
    Index2T = IndexTemp % (2 * GP.FluidsResolution)  # 进到周期
    Index1T = IndexTemp % (GP.FluidsResolution)  # 进到周期
    IndexArea = Index2T // GP.FluidsResolution  # 判断升降
    return Index2T - 2 * IndexArea * Index1T


@Ti.kernel
def CompactPoissonConvX(
        InputBuffer: Ti.template(), OutputBuffer: Ti.template(),
        PoissonWeights: Ti.template(), PoissonVectors: Ti.template(), Rank: int, KernelSize: int):
    for i, j in OutputBuffer:
        for k in range(KernelSize):  # KernelSize
            OutputBuffer[i, j] += \
                PoissonVectors[Rank, k] * PoissonWeights[Rank] * \
                InputBuffer[ \
                    Reflection(i + k - GP.KernelRadius), \
                    j]


@Ti.kernel
def CompactPoissonConvY(
        InputBuffer: Ti.template(), OutputBuffer: Ti.template(),
        PoissonVectors: Ti.template(), Rank: int, KernelSize: int):
    for i, j in OutputBuffer:
        for k in range(KernelSize):  # KernelSize
            OutputBuffer[i, j] += \
                PoissonVectors[Rank, k] * \
                InputBuffer[ \
                    i, \
                    Reflection(j + k - GP.KernelRadius)]


@Ti.kernel
def CompactPoissonAddTo(InputBuffer: Ti.template(), OutputBuffer: Ti.template()):
    for i, j in OutputBuffer:
        OutputBuffer[i, j] += InputBuffer[i, j]


@Ti.kernel
def CompactPoisson(
        InputBuffer: Ti.template(),
        OutputBuffer: Ti.template(),
        TempBuffer1: Ti.template(),
        TempBuffer2: Ti.template(),
        PoissonWeights: Ti.template(),
        PoissonVectors: Ti.template(),
        Rank: int,
        KernelSize: int):
    # CompactPoissonConvX(InputBuffer, TempBuffer1, PoissonWeights, PoissonVectors, Rank, KernelSize)
    # CompactPoissonConvY(TempBuffer1, TempBuffer2, PoissonVectors, Rank, KernelSize)
    for i, j in OutputBuffer:
        for k in range(KernelSize):  # KernelSize
            TempBuffer1[i, j] = TempBuffer1[i, j] + \
                                PoissonVectors[Rank, k] * PoissonWeights[Rank] * \
                                InputBuffer[ \
                                    Reflection(i + k - GP.KernelRadius), \
                                    j]
    for i, j in OutputBuffer:
        for k in range(KernelSize):  # KernelSize
            TempBuffer2[i, j] = TempBuffer2[i, j] + \
                                PoissonVectors[Rank, k] * \
                                TempBuffer1[ \
                                    i, \
                                    Reflection(j + k - GP.KernelRadius)]
    for i, j in OutputBuffer:
        OutputBuffer[i, j] += TempBuffer1[i, j]


def GeneralJacobiSolver(JacobiKernel, Buffer, *AppendentData):
    if JacobiKernel == Poisson:
        for _ in range(GP.JacobiIters):
            JacobiKernel(Buffer.cur, Buffer.nxt, AppendentData[0])
            Buffer.swap()
    elif JacobiKernel == Diffuse:
        for _ in range(GP.JacobiIters):
            JacobiKernel(Buffer.cur, Buffer.nxt)
            Buffer.swap()
    elif JacobiKernel == ScalarDiffuse:
        for _ in range(GP.JacobiIters):
            JacobiKernel(Buffer.cur, Buffer.nxt, AppendentData[0])
            Buffer.swap()
    elif JacobiKernel == CompactPoisson:
        # AppendentData[0] B
        # AppendentData[1] Inital Value
        # AppendentData[2] Output Buffer
        # AppendentData[3, 4] Temp
        # AppendentData[5, 6] Compact Poisson Weights/Vectors
        # AppendentData[7, 8] Initial Compact Poisson Weights/Vectors
        AppendentData[2].fill(0)
        # print(AppendentData[0])
        # print(AppendentData[1])
        # print(AppendentData[2])
        # print(AppendentData[3])
        # print(AppendentData[4])
        # print(AppendentData[5])
        # print(AppendentData[6])
        # print(AppendentData[7])
        # print(AppendentData[8])
        if (AppendentData[7] != None):
            for Rank in range(GP.MaxRank):
                AppendentData[3].fill(0)
                AppendentData[4].fill(0)
                CompactPoissonConvX(
                    AppendentData[1], AppendentData[3],
                    AppendentData[7], AppendentData[8], Rank, GP.KernelSize)
                CompactPoissonConvY(
                    AppendentData[3], AppendentData[4],
                    AppendentData[8], Rank, GP.KernelSize)
                CompactPoissonAddTo(
                    AppendentData[4], AppendentData[2])

        for Rank in range(GP.MaxRank):
            AppendentData[3].fill(0)
            AppendentData[4].fill(0)
            CompactPoissonConvX(
                AppendentData[0], AppendentData[3],
                AppendentData[5], AppendentData[6], Rank, GP.KernelSize)
            CompactPoissonConvY(
                AppendentData[3], AppendentData[4],
                AppendentData[6], Rank, GP.KernelSize)
            CompactPoissonAddTo(
                AppendentData[4], AppendentData[2])
